import numpy as np
import matplotlib.pyplot as plt


Population_Size = 100
Iteration_Number = 200
Cross_Rate = 0.8
Mutation_Rate = 0.003
Dna_Size = 10
X_Range=[0,5]


def F(x):
    '''
    目标函数，需要被优化的函数
    :param x:
    :return:
    '''
    return np.sin(10 * x) * x + np.cos(2 * x) * x

def CrossOver(Parent,PopSpace):
    '''
    交叉DNA，我们直接在种群里面选择一个交配
    然后就生出孩子了
    :param parent:
    :param PopSpace:
    :return:
    '''
    if(np.random.rand()) < Cross_Rate:

        cross_place = np.random.randint(0, 2, size=Dna_Size).astype(np.bool)
        cross_one = np.random.randint(0, Population_Size, size=1) #选择一位男/女士交配
        Parent[cross_place] = PopSpace[cross_one,cross_place]

    return Parent

def Mutate(Child):
    '''
    变异
    :param Child:
    :return:
    '''
    for point in range(Dna_Size):
        if np.random.rand() < Mutation_Rate:
            Child[point] = 1 if Child[point] == 0 else 0
    return Child


def TranslateDNA(PopSpace):
    '''
    把二进制转化为十进制方便计算
    :param PopSpace:
    :return:
    '''
    return PopSpace.dot(2 ** np.arange(Dna_Size)[::-1]) / float(2 ** Dna_Size - 1) * X_Range[1]

def Fitness(pred):
    '''
    这个其实是对我们得到的F(x)进行换算，其实就是选择的时候
    的概率，我们需要处理负数，因为概率不能为负数呀
    pred 这是一个二维矩阵
    :param pred:
    :return:
    '''

    return pred + 1e-3 - np.min(pred)
    # return (np.max(pred) - pred)+1e-3 求取最小值

def Select(PopSpace,Fitness):
    '''
    选择
    :param PopSpace:
    :param Fitness:
    :return:
    '''
    '''
    这里注意的是，我们先按照权重去选择我们的优良个体，所以我们这里选择的时候允许重复的元素出现
    之后我们就可以去掉这些重复的元素，这样才能实现保留良种去除劣种。100--》70（假设有30个重复）
    如果不允许重复的话，那你相当于没有筛选
    '''
    Better_Ones = np.random.choice(np.arange(Population_Size), size=Population_Size, replace=True,
                           p=Fitness / Fitness.sum())
    # np.unique(Better_Ones) #如果我要去掉不良基因的话，但是为了多样性，还是保留吧

    
    return PopSpace[Better_Ones]


if __name__ == '__main__':
    PopSpace = np.random.randint(2, size=(Population_Size, Dna_Size))  # initialize the PopSpace DNA

    plt.ion()
    x = np.linspace(X_Range, 200)
    x_show = np.linspace(0,5,100)
    y_show = F(x_show)
    plt.plot(x_show, y_show)
    plt.xticks([0,5])
    plt.yticks([0,5])

    for _ in range(Iteration_Number):
        F_values = F(TranslateDNA(PopSpace))

        # something about plotting
        if 'sca' in globals():
            sca.remove()
        sca = plt.scatter(TranslateDNA(PopSpace), F_values, s=200, lw=0, c='red', alpha=0.5)
        plt.pause(0.05)

        # GA part (evolution)
        fitness = Fitness(F_values)
        print("Most fitted DNA: ", PopSpace[np.argmax(fitness)])
        PopSpace = Select(PopSpace, fitness)
        PopSpace_copy = PopSpace.copy()
        for parent in PopSpace:

            child = CrossOver(parent, PopSpace_copy)
            child = Mutate(child)

            parent[:] = child

    plt.ioff()
    plt.show()
